import torch
from torch.utils.data import Dataset
from torchvision.transforms import transforms
import cv2
import numpy as np
import random
import os

class Transform(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 图片输入大小
        self.size = 640

        self.transform = transforms.Compose([
            transforms.ToTensor()
        ])

    def forward(self, img):
        h, w = img.shape[0:2]
        long_border = max(h, w)

        pad_img = np.full((long_border, long_border, 3), fill_value=114, dtype=np.uint8)
        pad_img[int((long_border - h)/2):h+int((long_border - h)/2), int((long_border - w)/2):w+int((long_border - w)/2), :] = img
        pad_img = cv2.resize(pad_img, (self.size, self.size))
        pad_img = self.transform(pad_img)

        return pad_img

class MyDataset(Dataset):
    def __init__(self, img_rootPath, mode='train', scale=0.9):
        super().__init__()
        
        # 数据集划分
        if mode == 'train':
            data_dir = list(os.listdir(img_rootPath))
            data_dir = data_dir[:int(len(data_dir)*scale)]
        elif mode == 'test':
            data_dir = list(os.listdir(img_rootPath))
            data_dir = data_dir[int(len(data_dir)*(scale)):]

        self.img_path_arr = []
        for img_name in data_dir:
            self.img_path_arr.append(os.path.join(img_rootPath, img_name))

        '''
            cv2.FONT_ITALIC：斜体字的标志
            cv2.FONT_HERSHEY_PLAIN：小尺寸无衬线字体
            cv2.FONT_HERSHEY_SIMPLEX：正常大小的无衬线字体
            cv2.FONT_HERSHEY_DUPLEX：正常大小的无衬线字体（比FONT_HERSHEY_SIMPLEX更复杂）
            cv2.FONT_HERSHEY_COMPLEX：正常大小的衬线字体
            cv2.FONT_HERSHEY_TRIPLEX：正常大小的衬线字体（比FONT_HERSHEY_COMPLEX更复杂）
            cv2.FONT_HERSHEY_SCRIPT_SIMPLEX：手写体字体
            cv2.FONT_HERSHEY_SCRIPT_COMPLEX（比FONT_HERSHEY_SCRIPT_SIMPLEX的更复杂）

        '''
        self.fontFace_list = [
            cv2.FONT_ITALIC,
            cv2.FONT_HERSHEY_PLAIN,
            cv2.FONT_HERSHEY_SIMPLEX,
            cv2.FONT_HERSHEY_DUPLEX,
            cv2.FONT_HERSHEY_COMPLEX,
            cv2.FONT_HERSHEY_TRIPLEX,
            cv2.FONT_HERSHEY_SCRIPT_SIMPLEX,
            cv2.FONT_HERSHEY_SCRIPT_COMPLEX
        ]

        self.to_tensor = Transform()
    
    def __getitem__(self, item):
        # 读取图片
        img_path = self.img_path_arr[item]
        img = cv2.imread(img_path)

        # 生成水印
        mask = np.zeros(shape=img.shape, dtype=np.uint8)
        for i in range(random.randint(2, 12)):
            text = ''
            for _ in range(11):
                text += str(random.randint(0, 9))
            fontFace = self.fontFace_list[random.randint(0, len(self.fontFace_list)-1)]
            fontScale = random.random() * 1.5 + 0.5
            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
            thickness = random.randint(1, 3)
            x, y = int(mask.shape[1]*random.random()), int(mask.shape[0]*random.random())

            # text_size = cv2.getTextSize(text, fontFace, fontScale, thickness)
            # w, h = text_size[0]

            cv2.putText(mask, text, (x, y), fontFace, fontScale, color, thickness)
            # cv2.rectangle(img, (x, y), (x+w, y-h), (0, 0, 255), 1)

        # 合成水印图片
        t = random.randint(0, 9)
        X_img = cv2.addWeighted(img, 1, mask, (t+1)*0.1, 0)
        traget_img = cv2.addWeighted(img, 1, mask, (t)*0.1, 0)

        # 转Tensor并归一化
        X_img = self.to_tensor(X_img) / 255.
        traget_img = self.to_tensor(traget_img) / 255.

        return X_img, traget_img

    def __len__(self):
        return len(self.img_path_arr)


if __name__ == '__main__':
    dataset = MyDataset(r'D:\VOCtrainval_11-May-2012\JPEGImages')
    print(dataset[5])